# -*- coding: utf-8 -*-
from __future__ import print_function
from triton import ARCH, TritonContext, Instruction, MODE, MemoryAccess, CPUSIZE
from triton import *
import os
import sys


EBP_ADDR = 0x100000
# 存放参数的地址
ARG_ADDR = 0x200000

Triton = TritonContext()
Triton.setArchitecture(ARCH.X86)

def init_machine():
    Triton.concretizeAllMemory()
    Triton.concretizeAllRegister()
    Triton.clearPathConstraints()
    Triton.setConcreteRegisterValue(Triton.registers.ebp, EBP_ADDR)

    # 设置栈
    Triton.setConcreteRegisterValue(Triton.registers.ebp, EBP_ADDR)
    Triton.setConcreteRegisterValue(Triton.registers.esp, EBP_ADDR -  0x2000)

    for i in range(2):
        Triton.setConcreteMemoryValue(MemoryAccess(EBP_ADDR - 0x14 + i * 4, CPUSIZE.DWORD), 5)

# 加载 elf 文件到内存
def loadBinary(path):
    import lief
    binary = lief.parse(path)
    phdrs = binary.segments
    for phdr in phdrs:
        size = phdr.physical_size
        vaddr = phdr.virtual_address
        print('[+] Loading 0x%06x - 0x%06x' % (vaddr, vaddr+size))
        Triton.setConcreteMemoryAreaValue(vaddr, phdr.content)
    return

def crack():
    i = 1
    Triton.setConcreteMemoryValue(MemoryAccess(EBP_ADDR - 0x14, CPUSIZE.DWORD), i)
    pc = 0x8048CED
    while pc:

        # x86 指令集的字节码的最大长度为 15
        opcode = Triton.getConcreteMemoryAreaValue(pc, 16)
        instruction = Instruction()
        instruction.setOpcode(opcode)
        instruction.setAddress(pc)
        Triton.processing(instruction)

        if instruction.getAddress() == 0x08048D01:
            print("solve!  answer: %d" %(i))
            break


        if instruction.getAddress() == 0x8048D07:
            pc = 0x8048CED
            i += 1
            # 重置运行时 
            init_machine()
            # 再次设置参数
            Triton.setConcreteMemoryValue(MemoryAccess(EBP_ADDR - 0x14, CPUSIZE.DWORD), i)
            continue
        pc = Triton.getConcreteRegisterValue(Triton.registers.eip)
    print('[+] Emulation done.')    


if __name__ == '__main__':
    init_machine()
    loadBinary(os.path.join(os.path.dirname(__file__), 'bomb'))
    crack()
    sys.exit(0)
